package io.dy.aspect;

import io.dy.anno.RateLimiter;
import io.dy.utils.IpUtil;
import lombok.RequiredArgsConstructor;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import java.lang.reflect.Method;
import java.util.Collections;

/**
 * @author : dmj
 * @date: 2024/6/25 4:04 下午
 * @description:
 */
@Aspect
@Component
@RequiredArgsConstructor
public class RateLimiterAspect {


    private final StringRedisTemplate redisTemplate;

    private static final DefaultRedisScript<Long> LIMIT_SCRIPT;

    static {
        LIMIT_SCRIPT = new DefaultRedisScript<>();
        LIMIT_SCRIPT.setLocation(new ClassPathResource("lua/rateLimiter.lua"));
        LIMIT_SCRIPT.setResultType(Long.class);
    }

    /**
     * 前置处理
     * @param point
     * @param rateLimiter
     * @throws Throwable
     */
    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter){

        //1、获取限流次数、时间
        int count = rateLimiter.count(), time = rateLimiter.time();

        //2、获取对应key
        String key = assembleKey(point, rateLimiter);

        //3、获取是否超过访问次数标识
        try{
            Long flag = redisTemplate.execute(LIMIT_SCRIPT, Collections.singletonList(key),  count + "", time + "");
            if(flag == null || flag == 0){
                throw new RuntimeException("访问过于频繁，请稍候再试");
            }
        }catch (RuntimeException e){
            throw e;
        }catch (Exception e){
            throw new RuntimeException("服务器限流异常，请稍候再试");
        }
    }

    /**
     * 组装key
     *  1、LimitType = DEFAULT 表示在指定时间内记录所有用户对该接口方法的请求方式 全局型 key = prefixKey + 类名 + ":" + 方法名
     *
     *  2、LimitType = IP 表示根据IP地址记录访问次数  key = prefixKey + 类名 + ":" + 方法名 + ":" + IP
     *
     *  3、LimitType = USER 表示根据用户id记录访问次数 key = prefixKey + 类名 + ":" + 方法名 + ":" + userId
     *
     * @param point
     * @param rateLimiter
     * @return
     */
    private String assembleKey(JoinPoint point, RateLimiter rateLimiter) {
        //1、组装共性部分
        StringBuffer buffer = new StringBuffer(rateLimiter.prefixKey());
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        buffer.append(":").append(targetClass.getName()).append(":").append(method.getName());

        //2、根据类型拼接
        switch(rateLimiter.limitType()){
            case IP -> {
                //2.1 获取访问接口IP
                String ipAddress = IpUtil.getIpAddress();
                buffer.append(":").append(ipAddress);
            }

            case USER -> buffer.append(":").append("10");
        }
        return buffer.toString();
    }
}
